from te import tik

NAME="train_mmad_model"


def exec_op(a, b, r, d, output, kernel_name):
    tik_instance = tik.Tik()

    input_a = tik_instance.Tensor("float16", a.get("shape"), name="a",  scope=tik.scope_gm)
    input_b = tik_instance.Tensor("float16", b.get("shape"), name="b",  scope=tik.scope_gm)
    input_r = tik_instance.Tensor("int32", r.get("shape"), name="r",  scope=tik.scope_gm)
    input_d = tik_instance.Tensor("int32", d.get("shape"), name="d",  scope=tik.scope_gm)


    input_a_ub = tik_instance.Tensor("float16", a.get("shape"), name="a_ub",  scope=tik.scope_cbuf)
    input_b_ub = tik_instance.Tensor("float16", a.get("shape"), name="b_ub",  scope=tik.scope_cbuf)
    duration_ub = tik_instance.Tensor("int32", d.get("shape"), name="d_ub", scope=tik.scope_ubuf)
    output_ub = tik_instance.Tensor("float32", output.get("shape"), name="out_ub",  scope=tik.scope_ubuf)
    output_cb = tik_instance.Tensor("float32", output.get("shape"), name="out_ub",  scope=tik.scope_cbuf_out)

    output_gm = tik_instance.Tensor("float32", output.get("shape"), name="out_ub",  scope=tik.scope_gm)

    tik_instance.data_move(input_a_ub, input_a, 0, 1, 32, 0, 0)
    tik_instance.data_move(input_b_ub, input_b, 0, 1, 32, 0, 0)
    tik_instance.data_move(duration_ub, input_d, 0, 1, 1, 0, 0)
    # duration_times 表示重复计算多少次
    duration_times = tik_instance.Scalar(dtype="int32")
    duration_times.set_as(duration_ub[0,0])

    with tik_instance.for_range(0, duration_times, thread_num=1):
        tik_instance.matmul(output_cb, input_a_ub, input_b_ub, 256, 32, 128)


    tik_instance.data_move(output_ub, output_cb, 0, 1, 1, 0, 0)
    tik_instance.data_move(output_gm, output_ub, 0, 1, 1, 0, 0)

    tik_instance.BuildCCE(kernel_name,
                          inputs=[input_a, input_b, input_r, input_d],
                          outputs=[output_gm])

#maxinum 112 KB
M = 256
N = 32
K = 128
ROW = 1
COL = 8
INPUT_A = {"shape": [M, N], "dtype": "float16"}
INPUT_B = {"shape": [N, K], "dtype": "float16"}
INPUT_R = {"shape": [ROW, COL], "dtype": "int32"}
INPUT_D = {"shape": [ROW, COL], "dtype": "int32"}
OUTPUT = {"shape": [M, K], "dtype": "float32"}

exec_op(INPUT_A, INPUT_B, INPUT_R, INPUT_D, OUTPUT, NAME)
